#include "MPU6500.hpp"
#include "I2CDevice.hpp"
#include "Timer.hpp"
#include <cassert>
#include <iostream>
#include "Model.hpp"
#include "ThreadPool.hpp"
#include "MultiBuffer.hpp"
#include <mutex>
#include <iomanip>

#include <functional>
#include <signal.h>

std::function<void()> appHandler;

void signalHandler(int signum) {
    std::cout << "Interrupt signal (" << signum << ") received.\n";
    appHandler();
    exit(signum);
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << "i2c_device_path" << std::endl;
        return 1;
    }

    float ax, ay, az, gx, gy, gz;
    // float buffer[128][6] = {0};
    // int buffer_index = 0;
    MultiBufferManager<float> Ax(128, 4), Ay(128, 4), Az(128, 4), Gx(128, 4), Gy(128, 4), Gz(128, 4);
    std::mutex buffer_mutex; // Mutex for buffer protection

    try {
        I2CDevice i2c(argv[1], 0x68);
        MPU6500 mpu(i2c);

        Timer timer(20, [&](uint64_t now_us) {
            std::lock_guard<std::mutex> lock(buffer_mutex); // Lock before writing
            mpu.readSensor(ax, ay, az, gx, gy, gz);
        });

        appHandler = [&]() {
            timer.stop();
            i2c.close();          
        };
        signal(SIGINT, signalHandler); // Register signal handler for Ctrl+C

        mpu.initialize();
        mpu.calibrate();

        Model model("Model.mnn");
        
        const float mean[6] = {-0.00066533f, -0.00031751f, -0.00020395f, -0.0034113f, -0.00017161f, 0.00089603f};
        const float std[6] = {0.14196804f, 0.08511718f, 0.07230477f, 0.26109824f, 0.27211866f, 0.17970076f};
        
        ThreadPool pool(2); // Create a thread pool with 2 worker threads

        timer.start();

        while (true) {
            timer.wait();  // 等待数据采集完成
            {
                std::lock_guard<std::mutex> lock(buffer_mutex); // Lock before writing
                // buffer[buffer_index][0] = (ax - mean[0]) / std[0];
                // buffer[buffer_index][1] = (ay - mean[1]) / std[1];
                // buffer[buffer_index][2] = (az - mean[2]) / std[2];
                // buffer[buffer_index][3] = (gx - mean[3]) / std[3];
                // buffer[buffer_index][4] = (gy - mean[4]) / std[4];
                // buffer[buffer_index][5] = (gz - mean[5]) / std[5];
                // buffer_index = (buffer_index + 1) % 128;
                assert(Ax.write((ax - mean[0]) / std[0]));
                assert(Ay.write((ay - mean[1]) / std[1]));
                assert(Az.write((az - mean[2]) / std[2]));
                assert(Gx.write((gx - mean[3]) / std[3]));
                assert(Gy.write((gy - mean[4]) / std[4]));
                assert(Gz.write((gz - mean[5]) / std[5]));
            }
            // if (buffer_index % 128 == 0) {
            if (Ax.availableBuffers() > 0) { // Check if there's at least one full buffer
                float buffer[128][6]; // Temporary buffer to hold the data
                pool.enqueue([&]() {
                    {
                        std::lock_guard<std::mutex> lock(buffer_mutex);
                        Ax.read(&buffer[0][0], 6);
                        Ay.read(&buffer[0][1], 6);
                        Az.read(&buffer[0][2], 6);
                        Gx.read(&buffer[0][3], 6);
                        Gy.read(&buffer[0][4], 6);
                        Gz.read(&buffer[0][5], 6);
                    }
                    model.fillInputTensor(buffer); // Fill the input tensor with the buffer data
                    std::array<float, 6> outputs;
                    uint64_t prediction_start = Timer::getCurrentTimeMicroseconds();
                    model.predict(outputs);
                    uint64_t prediction_end = Timer::getCurrentTimeMicroseconds();
                    size_t prediction = Model::argmax(outputs);
                    // Print the prediction and duration
                    std::cout << "Prediction: " << model.getLabel(prediction) << ", Spending Time: " << (prediction_end - prediction_start) / 1000 << " ms, Probability: " << '[' << std::fixed << std::setprecision(4) << outputs[0] << ", " << outputs[1] << ", " << outputs[2] << ", " << outputs[3] << ", " << outputs[4] << ", " << outputs[5] << ']' << std::endl;
                });
            }
        }
    } catch (const std::exception& ex) {
        std::cerr << "Exception: " << ex.what() << std::endl;
    }

    return 0;
}
